import cv2
import numpy as np
import threading
from typing import Optional, Tuple

class NonBlockingImageWindow:
    """
    非阻塞图像窗口（仅支持 fmt='path' 或 'rgb'）。

    主要特性：
    - 默认窗口创建在屏幕中心；若提供 position=(x, y) 则按坐标显示
    - 新增：fullscreen 模式，窗口全屏显示，分辨率=屏幕分辨率，忽略 size/position 及移动类操作  # NEW: fullscreen
    - start(): 启动显示线程（__init__ 默认已启动）
    - update(img, fmt): 更新图像内容，fmt ∈ {'path','rgb'}
    - stop(timeout=None): 关闭窗口并等待显示线程退出（自动 join）
    - move_to(x, y): 运行中移动窗口到指定坐标（非全屏时有效）
    - center(): 运行中重新居中窗口（非全屏时有效）
    - set_fullscreen(on): 运行时切换全屏/窗口模式                                 # NEW: fullscreen
    """

    def __init__(self,
                 title: str = "Image",
                 size: Optional[Tuple[int, int]] = None,   # (w, h)（非全屏时生效）
                 fps: int = 30,
                 auto_start: bool = True,
                 position: Optional[Tuple[int, int]] = None,  # 非全屏时生效
                 fullscreen: bool = False  # NEW: fullscreen
                 ):
        self.title = title
        self.size = size
        self.delay = max(1, int(1000 / max(1, fps)))
        self._img = None
        self._lock = threading.Lock()
        self._running = threading.Event()
        self._thread: Optional[threading.Thread] = None

        self._init_position = position                    # 初始位置（None 表示居中；仅非全屏）
        self._pending_move: Optional[Tuple[int, int]] = None
        self._pending_center: bool = False

        self._fullscreen = bool(fullscreen)               # NEW: fullscreen 当前状态
        self._pending_fullscreen: Optional[bool] = None   # NEW: 运行时切换请求
        self._screen_wh: Optional[Tuple[int, int]] = None # NEW: 屏幕分辨率缓存

        if auto_start:
            self.start()

    # ------------ 公共方法 ------------
    def start(self):
        if self._thread and self._thread.is_alive():
            return
        self._running.set()
        self._thread = threading.Thread(target=self._loop, daemon=True)
        self._thread.start()

    def update(self, img, fmt: str):
        """fmt 只能是 'path' 或 'rgb'"""
        arr = self._to_bgr(img, fmt)
        with self._lock:
            self._img = arr

    def stop(self, timeout: Optional[float] = None):
        """关闭窗口并等待显示线程退出（自动 join）。"""
        self._running.clear()
        if self._thread and threading.current_thread() is not self._thread:
            self._thread.join(timeout)
        try:
            cv2.destroyWindow(self.title)
        except Exception:
            pass

    def join(self, timeout: Optional[float] = None):
        if self._thread and threading.current_thread() is not self._thread:
            self._thread.join(timeout)

    def move_to(self, x: int, y: int):
        """
        将窗口移动到 (x, y)（窗口左上角）。
        全屏模式下无效（被忽略）。                                  # NEW: fullscreen
        """
        with self._lock:
            if self._fullscreen:
                return
            self._pending_move = (int(x), int(y))
            self._pending_center = False

    def center(self):
        """
        使窗口重新居中。
        全屏模式下无效（被忽略）。                                  # NEW: fullscreen
        """
        with self._lock:
            if self._fullscreen:
                return
            self._pending_center = True
            self._pending_move = None

    def set_fullscreen(self, on: bool):
        """
        运行中切换全屏/窗口模式。
        全屏开启后：忽略 size/position/move_to/center，分辨率固定为屏幕分辨率。
        """
        with self._lock:
            self._pending_fullscreen = bool(on)

    # ------------ 内部实现 ------------
    def _loop(self):
        cv2.namedWindow(self.title, cv2.WINDOW_NORMAL)

        # 屏幕分辨率
        screen_w, screen_h = self._get_screen_size_with_tk()
        if screen_w is not None and screen_h is not None:
            self._screen_wh = (int(screen_w), int(screen_h))

        # 初始化窗口显示模式
        self._apply_mode(first_time=True)

        try:
            while self._running.is_set():
                frame = None
                with self._lock:
                    if self._img is not None:
                        frame = self._img

                    # 取出一次性意图
                    pending_move = self._pending_move
                    pending_center = self._pending_center
                    pending_fullscreen = self._pending_fullscreen

                    # 清除意图
                    self._pending_move = None
                    self._pending_center = False
                    self._pending_fullscreen = None

                # 切换全屏/窗口模式（优先执行）
                if pending_fullscreen is not None:
                    self._fullscreen = pending_fullscreen
                    self._apply_mode()

                    # 切到窗口模式后，可恢复一次位置逻辑（若此前请求了移动/居中）
                    if (not self._fullscreen) and pending_move is not None:
                        try:
                            cv2.moveWindow(self.title, int(pending_move[0]), int(pending_move[1]))
                        except Exception:
                            pass
                    elif (not self._fullscreen) and pending_center:
                        try:
                            self._center_window()
                        except Exception:
                            pass

                # 渲染（全屏时：将帧拉伸到屏幕分辨率）
                if frame is not None:
                    if self._fullscreen and self._screen_wh is not None:
                        sw, sh = self._screen_wh
                        if frame.shape[1] != sw or frame.shape[0] != sh:
                            frame = cv2.resize(frame, (sw, sh), interpolation=cv2.INTER_LINEAR)
                    cv2.imshow(self.title, frame)

                # 仅在非全屏时响应位置调整
                if (not self._fullscreen) and (pending_move is not None):
                    try:
                        cv2.moveWindow(self.title, int(pending_move[0]), int(pending_move[1]))
                    except Exception:
                        pass
                elif (not self._fullscreen) and pending_center:
                    try:
                        self._center_window()
                    except Exception:
                        pass

                k = cv2.waitKey(self.delay)
                if k == 27:  # ESC
                    self._running.clear()
                    break
        finally:
            try:
                cv2.destroyWindow(self.title)
            except Exception:
                pass

    def _apply_mode(self, first_time: bool = False):
        """
        根据 self._fullscreen 应用窗口模式/全屏模式的窗口属性。
        """
        # 若支持屏幕分辨率
        sw, sh = (self._screen_wh or (None, None))

        if self._fullscreen:
            # 全屏：强制全屏属性 & 尺寸=屏幕分辨率
            try:
                cv2.setWindowProperty(self.title, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)  # 全屏
            except Exception:
                pass
            if sw and sh:
                try:
                    cv2.resizeWindow(self.title, sw, sh)  # 强制屏幕分辨率
                except Exception:
                    pass
            # 全屏下不处理位置
        else:
            # 窗口模式：退出全屏并应用 size/position/center
            try:
                cv2.setWindowProperty(self.title, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL)
            except Exception:
                pass

            # 若指定了 size，则调整窗口大小
            if self.size and self.size[0] > 0 and self.size[1] > 0:
                try:
                    cv2.resizeWindow(self.title, self.size[0], self.size[1])
                except Exception:
                    pass

            # 首次启动：优先使用用户指定坐标；否则居中一次
            if first_time:
                try:
                    if self._init_position is not None:
                        x, y = self._init_position
                        cv2.moveWindow(self.title, int(x), int(y))
                    else:
                        self._center_window()
                except Exception:
                    pass

    def _center_window(self):
        """将窗口移动到屏幕中心；依赖 tkinter 获取屏幕分辨率（失败则静默跳过）。"""
        screen_w, screen_h = self._screen_wh or self._get_screen_size_with_tk()
        if not screen_w or not screen_h:
            return

        # 估计窗口大小：优先用用户给定的 size；否则使用默认 640x480
        if self.size and self.size[0] > 0 and self.size[1] > 0:
            win_w, win_h = self.size
        else:
            win_w, win_h = 640, 480

        pos_x = max(0, (screen_w - win_w) // 2)
        pos_y = max(0, (screen_h - win_h) // 2)
        cv2.moveWindow(self.title, int(pos_x), int(pos_y))

    @staticmethod
    def _get_screen_size_with_tk():
        """用 tkinter 获取屏幕分辨率；失败返回 (None, None)。"""
        try:
            import tkinter as tk
            root = tk.Tk()
            root.withdraw()
            w = root.winfo_screenwidth()
            h = root.winfo_screenheight()
            root.destroy()
            return int(w), int(h)
        except Exception:
            return None, None

    def _to_bgr(self, img, fmt: str) -> np.ndarray:
        fmt = fmt.lower()
        if fmt == "path":
            arr = cv2.imread(str(img), cv2.IMREAD_UNCHANGED)
            if arr is None:
                raise ValueError(f"读取文件失败: {img}")
            return arr  # imread 已是 BGR/BGRA/Gray

        if fmt == "rgb":
            if not isinstance(img, np.ndarray):
                raise TypeError("fmt='rgb' 时 img 必须是 numpy.ndarray")
            if img.ndim != 3 or img.shape[2] != 3:
                raise ValueError("fmt='rgb' 需要形状 (H, W, 3) 的 RGB 图像")
            return cv2.cvtColor(self._normalize_to_uint8(img), cv2.COLOR_RGB2BGR)

        raise ValueError(f"不支持的 fmt: {fmt}（仅支持 'path' 和 'rgb'）")

    def _normalize_to_uint8(self, arr: np.ndarray) -> np.ndarray:
        if arr.dtype == np.uint8:
            return arr
        a = arr.astype(np.float32)
        a_min = float(np.nanmin(a))
        a_max = float(np.nanmax(a))
        if a_max == a_min:
            return np.zeros_like(a, dtype=np.uint8)
        a = (a - a_min) / (a_max - a_min) * 255.0
        return np.clip(a, 0, 255).astype(np.uint8)

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.stop()
